from pycocotools.coco import COCO
import requests
import os
import json
from PIL import Image
import torchvision.transforms as transforms

# 初始化COCO API
dataDir = './annotations_trainval2017'  # 数据集所在目录
annFile = f'{dataDir}/annotations/instances_train2017.json'
coco = COCO(annFile)

# 获取“人”类别的ID
catIds = coco.getCatIds(catNms=['person'])
imgIds = coco.getImgIds(catIds=catIds)

# 下载选定的图像，限制为2000张
download_limit = 500
downloaded_count = 0
annotations = []
images = []
# 创建保存图像的目录
os.makedirs('downloaded_images5/train_img', exist_ok=True)
os.makedirs('downloaded_images5/annotations_json', exist_ok=True)
os.makedirs('downloaded_images7/train_img', exist_ok=True)
# 定义数据增强的转换
data_transforms = transforms.Compose([
    transforms.Resize((256, 256)),  # 统一尺寸
    transforms.RandomRotation(degrees=15),  # 随机旋转
    transforms.RandomHorizontalFlip(),  # 随机水平翻转
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),  # 改变颜色
    transforms.ToTensor(),  # 转换为Tensor
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 归一化
])
for imgId in imgIds:
    if downloaded_count >= download_limit:
        break
    img = coco.loadImgs(imgId)[0]
    img_url = img['coco_url']# URL 指向图像的在线位置，可以用来下载图像数据。
    img_data = requests.get(img_url).content#使用 requests 库向 img_url 发送 HTTP GET 请求，以下载图像的内容
    # 保存原始图像
    original_img_path = os.path.join('./downloaded_images/train_img', img['file_name'])
    with open(original_img_path, 'wb') as handler:
        handler.write(img_data)
    # 图像处理
    image = Image.open(original_img_path).convert('RGB')
    augmented_image = data_transforms(image)
    # 检查是否为三维张量 (C, H, W)
    if augmented_image.ndimension() == 3:
        augmented_image = augmented_image.permute(1, 2, 0)  # 转换为 (H, W, C)
    # 根据需要缩放像素值
    augmented_image = (augmented_image*255).clamp(0, 225).byte().numpy()  # 缩放到 0-255
    # 转换为 PIL 图像
    augmented_image_pil = Image.fromarray(augmented_image)
    # 保存处理后的图像
    augmented_img_path = os.path.join('downloaded_images2/train_img', f'aug_{img["file_name"]}')
    augmented_image_pil.save(augmented_img_path)
    # 获取该图像的标注
    annIds = coco.getAnnIds(imgIds=img['id'], catIds=catIds)
    anns = coco.loadAnns(annIds)
    # 保存图像信息
    images.append({
        'id': img['id'],
        'file_name': f'aug_{img["file_name"]}',
        'width': 256,
        'height': 256
    })
    # 保存标注信息
    for idx, ann in enumerate(anns):
        # 将边界框转换为 COCO 格式 [x, y, width, height]
        x, y, w, h = ann['bbox']
        annotations.append({
            'id': downloaded_count * 100 + idx + 1,  # 为每个标注分配唯一ID
            'image_id': img['id'],
            'bbox': [x, y, w, h],  # 边界框
            'category_id': ann['category_id']  # 类别ID
        })
    downloaded_count += 1
# 构建符合COCO格式的字典
coco_format = {
    'images': images,
    'annotations': annotations,
    'categories': [{'id': catId, 'name': coco.loadCats(catId)[0]['name']} for catId in catIds]
}
# 保存标注到 JSON 文件
with open('downloaded_images2/annotations_json/annotations.json', 'w') as f:
    json.dump(coco_format, f)
print(f'Downloaded {downloaded_count} images and saved annotations.')
